# 加载环境变量
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

# 实例化模型本地代理
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
embeddings  = ZhipuAIEmbeddings(model="embedding-3")
llm = ChatZhipuAI(model=model,
                  temperature=0.9,              
    ) 

# 加载PDF文件内容
from langchain_community.document_loaders import PyPDFLoader
loader = PyPDFLoader("data/story.pdf")
documents = loader.load()

# 切分文档
from langchain.text_splitter import CharacterTextSplitter
text_splitter = CharacterTextSplitter(chunk_size=100, 
                                      chunk_overlap=50
    )
texts = text_splitter.split_documents(documents)
# print(texts)

# 向量化文档内容并存储于向量库中
# from langchain_community.vectorstores import Chroma
from langchain_chroma import Chroma
db = Chroma.from_documents(documents = texts, 
                           embedding = embeddings
    )

# 实例化QA检索器
from langchain.chains import RetrievalQA
retriever = db.as_retriever()
qa = RetrievalQA.from_chain_type(llm=llm, 
                                 chain_type="stuff", 
                                 retriever=retriever
    )

# 测试文档问答
query = "谁是医生？"
ret = qa.invoke(query)
print(f"Query:\n{ret['query']}")
print(f"Result:\n{ret['result']}\n")

query = "艾丽斯的现任男友是谁？"
ret = qa.invoke(query)
print(f"Query:\n{ret['query']}")
print(f"Result:\n{ret['result']}\n")